Feat: flux2 dev support - #9234

Frontend build is failing. pytests are failing.

Round 2 review at 75e9a64. First, thank you — all 10 round-1 findings were fixed correctly, and the dedup follow-ups came out very clean (details in my earlier comment). This pass covers the fix commits plus a fresh sweep of the whole branch.

4 blockers (🔴) and 6 nice-to-haves (🟡) as inline comments. The one to look at first is the v3→v4 persist migration: as written it wipes the entire params slice for every upgrading user — the fix is one line plus a test-fixture correction.

Cleanup notes (non-blocking, take or leave):

• configs/main.py:867 + :893-900 — with the new layout guard in place, the Flux2Transformer2DModel class-name entry and the root-config.json variant fallback ("loose transformer-only checkouts") are dead in practice; the comments now describe a layout the guard rejects. Worth dropping both and tightening the comments so nobody re-enables that path.
• No identification tests for the three Mistral encoder probes or the new pipeline layout guard — tests/backend/model_manager/configs/test_qwen3_encoder_config.py shows the cheap synthetic-dir pattern; this probe has been the most-revised code in the PR and would benefit most.
• load/model_loaders/mistral_encoder.py:177-189 — the three single-key GGUF metadata helpers are dead (they only call each other; production uses the batched reader). Safe to delete.
• GGUF header is still parsed twice per text-encoder load (once in gguf_sd_loader, once in _read_gguf_metadata_values) plus once for the tokenizer submodel — each pass re-walks the 131k-token vocab arrays (~1s each). Folding the metadata read into the gguf_sd_loader pass would fix it.
• configs/lora.py:105 — _FLUX2_VEC_IN_DIMS is a hand-maintained copy of flux2_variant.py's _VEC_IN_DIM values; exporting a FLUX2_VEC_IN_DIMS frozenset would close the last drift gap the shared module was created to prevent.
• load/model_loaders/mistral_encoder.py:799-814 — the missing-norm re-init loop is fully subsumed by _materialize_remaining_meta_tensors (same init, dtype, device, norm test); its only net effect is suppressing the materialize warning on the checkpoint path while the GGUF path warns for the identical condition. Deleting it deduplicates and unifies the diagnostics.
• _drop_quantization_metadata:352 — scale.repeat_interleave materializes a full weight-shape fp32 scale (3 transient copies vs 2); a reshape+broadcast multiply is a drop-in saving ~0.7-2.7GB on the largest tensors.
• parsing.test.tsx:34-113 — the fakeMainModel/modelRegistry fixtures and the it.each variant axis are dead (the merged handler never reads metadata.model; the registry key is never looked up), and the comment claims variant disambiguation that no longer exists. The two variant cases run the same test twice.

Checked and deliberately not raised: the Klein LoRA loaders' warn→raise change without a version bump (a dev-variant LoRA config cannot exist pre-PR, so no working workflow changes behavior); the llama.cpp→transformers key converter's replace chain (the trailing-dot patterns are collision-safe and unknown keys pass through unchanged, matching the z_image converter); GGUF reader mmap lifetime on Windows (the raw readers are transient locals; the long-lived-view problem WrappedGGUFReader solves doesn't apply).

🤖 Generated with Claude Code

I've re-verified the branch at 3e62aa8 against my second review. Thank you for the thorough fix commit (5026c02) — 8 of the 10 findings are correctly and completely fixed, including all the hard parts (the migration seed with a field-accurate test fixture, model-resolving guidance recall, the _reinit_inv_freq extraction validated against transformers 5.5.4, and the vision-tower drop). Two of the fixes need a follow-up, which is why this is another change request rather than an approval. Details below.

Requested changes

🔴 1. The new Klein-loader variant guard broke the VAE-only fallback (regression from the blocker-3 fix)

flux2_klein_model_loader.py calls _validate_diffusers_format on both call sites — the encoder-extraction path (:154) and the VAE-only extraction path (:132). The encoder-side rejection is exactly right, but the 32-channel AutoencoderKLFlux2 is shared between Klein and [dev], and the linear UI deliberately relies on that: buildFLUXGraph.ts falls back to any FLUX.2 diffusers pipeline when only the VAE is needed ("Fall back to any diffusers model if only the VAE is needed"), and isFlux2DiffusersMainModelConfig / readiness don't filter by variant.

Triggering sequence: install a Klein GGUF main + a standalone Qwen3 encoder + a FLUX.2 [dev] diffusers pipeline as the only FLUX.2 diffusers model (no standalone VAE) → readiness enables Invoke → the builder sets qwen3_source_model to the [dev] pipeline for VAE extraction → the new guard raises ValueError claiming the VAE is "incompatible with the Klein transformer" — which worked before this fix, and the message is factually wrong for the VAE.

Suggested fix: apply the variant guard only at the tokenizer/encoder extraction call sites ([dev] loader :153, Klein loader :154), not the VAE ones — or, if you prefer to reject cross-variant VAE sourcing on principle, the frontend readiness + source-selection fallback must be variant-filtered to match, so users get a disabled Invoke button with a reason instead of a runtime error.

🔴 2. The root-directory tokenizer step still fails offline for the exact layout blocker 4 described

The structural cleanup is all there (dead fallback gone, duplicated pre-try gone, ValueError in every except tuple, all three loaders share the helper). But the new step 3 (mistral_encoder.py:676-681) tries AutoProcessor only. With the pinned transformers 5.5.4, AutoProcessor.from_pretrained on a directory whose config.json is model_type: "mistral3" — which is precisely the BFL-style standalone-encoder layout this step targets — raises OSError: Can't load image processor … (it resolves mistral3 to a multimodal processor and requires preprocessor_config.json) before reaching any tokenizer fallback. The exception is swallowed and the ladder falls through to the HF fetch, which raises RuntimeError offline.

Reproduced empirically against the project venv with HF_HUB_OFFLINE=1: a root directory with a mistral3 config.json plus a valid tekken.json (or tokenizer.json + tokenizer_config.json) fails via AutoProcessor, while AutoTokenizer.from_pretrained on the same directory succeeds and returns a working backend. So the tokenizer is loadable — only the loader-class choice loses it. The step only succeeds today when the root config is text-only model_type: "mistral" or a full processor set including preprocessor_config.json is present.

Suggested fix: loop (AutoProcessor, AutoTokenizer) in steps 2 and 3, exactly as _load_tokenizer_from_hf already does (:616). Note if you do: probing showed AutoTokenizer can raise KeyError: 'special_tokens' on a directory containing only tekken.json, so the except tuple there needs KeyError too (or that escapes loudly).

Note on the two commits I just pushed (merge of main + migration v5 bump)

I've pushed 567cba9 (merge of latest main, which brings in the PiD decoder work — the flux2 PiD wiring is adapted into this branch's addFlux2Features helper and verified behavior-identical to main's) and d3c7bde. The second commit fixes a semantic conflict the textual merge couldn't see, and it wasn't your bug: main (after the PiD PR, 802de41) and this branch both bumped the params persist schema to _version: 4 with different new keys (main: pidMode/pidDecoderModel/gemma2EncoderModel/pidSteps, still with kleinVaeModel; this branch: flux2VaeModel + flux2DevMistralEncoderModel). After a plain merge, a v4 blob written by either parent is missing the other side's keys, skips every migration branch, fails zParamsState.parse(), and silently wipes the whole params slice — the same failure mode as blocker 1, reintroduced by the version collision.

The fix keeps main's v3→v4 verbatim and moves this branch's flux2 slot-merge + Mistral seed into a new v4→v5 step with conditional (?? null) seeding for both v4 shapes, with tests covering both. It also seeds the five wan* keys in v3→v4: main added them to the schema without a version bump while v6.13.x releases were still writing v3 blobs, so genuine released-build v3 blobs fail parse on those keys today (empirically confirmed — inherited from main, needs the same fix there).

Verified fixed (no action needed)

1. Migration seed (blocker 1) — flux2DevMistralEncoderModel = null seeded in the migration; the rebuilt fixture genuinely covers the regression (deleting the seed fails the test). ✅
2. Guidance for [dev] (blocker 2) — the parse gate resolves the image's own model and exempts variant === 'dev'; Klein still skips; slider renders for dev only via the same predicate the graph builder uses, so visibility and consumption can't drift; generate → metadata → recall → regenerate round-trips. Uninstalled-model dev images fall back to skipping guidance — reasonable fail-safe. ✅
3. Encoder-side variant guards (blocker 3) — the reported Klein-encoder-into-dev-transformer path (and mirror) is closed on every reachable route with clear ValueErrors, and the linear UI filters dev sources by variant. (VAE-path over-reach above is the only issue.) ✅
4. Tokenizer structure (blocker 4) — dead fallback and duplicate pre-try removed, ValueError caught, single helper used by all three loaders. (AutoProcessor-only gap above is the only issue.) ✅
5. Collection LoRA base check — ValueError before dedup/config fetch; FLUX.1 rejected; round-1 Klein-variant rejection intact; Klein side symmetric. ✅
6. Mistral slot clearing — cleared on base switch (with toast accounting) and intra-flux2 variant switch, including when the Klein encoder is unset; recall path covered via modelSelected. ✅
7. rope_theta — single _reinit_inv_freq with the rope_parameters/rope_scaling fallback; verified on transformers 5.5.4 that rope_theta set via MistralConfig(rope_theta=…) lands exactly where the fallback reads it. ✅
8. mistral-common pin — >=1.5.4,<2 in the right block, lock resolves 1.11.6, and both private-API usages (_special_tokens_reverse_vocab, Tekkenizer.encode(s, bos, eos)) verified present in-range with defensive fallbacks. ✅
9. Vision tower drop — setattr(model, …, None) removes the tower/projector from state_dict()/named_parameters(), so cache accounting and VRAM transfers exclude it; .language_model access and index-based extraction unaffected. Residual, non-blocking: the sibling full-pipeline path (flux.py Flux2DiffusersModel._load_model, used for diffusers mains and mistral_source_model) still loads and carries the ~0.8 GB tower — worth the same treatment in a follow-up. ✅
10. 40-layer docs — all four sources now agree on the operational contract (accepted, warned, indices 10/20/30, final norm kept for 40-layer / stripped for cow). Nit, non-blocking: starter_models.py says BFL's own 40-layer has weaker adherence due to index depth, while the loader warning implies only non-BFL 40-layer weights degrade — pick one story. ✅

Pre-existing, out of scope for this PR (noting for completeness): the workflow editor can still feed a Klein 9B-family source into a 4B main (no cross-Klein variant check on the source-extraction path — same wrong-width class, linear UI is protected); the Klein loader never validates its main model's variant.

🤖 Generated with Claude Code

Round 4 review (head c2b7778)

Thanks for the quick turnaround on both round-3 blockers. I re-verified each fix in code and with fresh-context adversarial passes (including mutation-testing your new tests — they do catch reverts). Summary:

• Round-3 blocker 1 (cross-variant guard also rejected VAE-only sourcing): fixed correctly. The split is right, every source-based encoder extraction goes through _validate_encoder_source, no VAE branch retains variant logic, and the linear UI can't wire a cross-variant source into an encoder slot in any main-format × standalone-VAE × standalone-encoder combination I enumerated.
• Round-3 blocker 2 (AutoProcessor-only local rungs): the offline gap itself is fixed — both local rungs now loop (AutoProcessor, AutoTokenizer). Unfortunately the widened AutoTokenizer rung introduces two new problems, one of which silently corrupts conditioning. Details below; both reproduced empirically against transformers 5.5.4 with real files, not hypothesized.
• The core_metadata mistral_encoder declaration + 2.2.0 bump looks good.

Housekeeping: I pushed two commits

• 7179712 — merge of current main. Conflicts: addRegions.ts (main refactored the posCond unions into a shared RegionalPositiveConditioning alias for Krea-2 regional prompting; I added flux2_dev_text_encoder to the alias and took main's shape), and schema.ts/openapi.json (regenerated from the merged backend rather than hand-merged).
• c2b7778 — the merge surfaced a semantic conflict git can't see: main's #9428 added a registry guard asserting every *_text_encoder node declares idle_gpu_offloadable=True, and flux2_dev_text_encoder (which neither parent knew about) failed it, so python-tests would have gone red. The flag alone would have been wrong — the marker's contract is that saved conditioning is CPU-backed, because the borrowed GPU's pool lock is released the moment the node returns. I mirrored the Klein encoder: detach().to("cpu") before save, marker added, version 1.0.1, plus the same output-device regression test Klein has. tests/app is fully green at head (2161 passed), as are the frontend suite (1736) and lints.

Blocker 1: AttributeError escapes the ladder and crashes the whole load

_TOKENIZER_LOAD_ERRORS = (OSError, EnvironmentError, ValueError, KeyError) (mistral_encoder.py:623) doesn't cover a failure mode of the very rung you added. In transformers 5.5.4, AutoTokenizer.from_pretrained on a directory whose tokenizer_config.json names a tokenizer_class unknown to 5.5.4 hits tokenizer_class_from_name(...) → None and raises AttributeError: 'NoneType' object has no attribute 'from_pretrained' (tokenization_auto.py:723 has no None-guard).

Triggering sequence (reproduced): encoder dir with config.json (model_type: "mistral3"), tekken.json, and a tokenizer_config.json containing "tokenizer_class": "MistralCommonTokenizer" — exactly what transformers 4.5x's MistralCommonTokenizer.save_pretrained writes. Rung 3's AutoProcessor raises OSError (caught), AutoTokenizer raises AttributeError → escapes _try_load_tokenizer_from_dir → _load_tokenizer_for_model crashes instead of falling through to the HF rung. Before this commit the same directory worked whenever rung 4 could succeed.

Given the ladder's whole contract is graceful fallthrough, I'd catch broadly in _try_load_tokenizer_from_dir (add AttributeError, or catch Exception and log) — an exotic exception from a probe rung should never be able to kill the load while a working fallback remains.

Blocker 2: root tekken.json + config.json now loads MistralCommonBackend, which tokenizes the template markers wrong — silently

The commit comment assumes tekken-carrying directories fail with KeyError: 'special_tokens'. That's only true when config.json is absent. With config.json (model_type: "mistral3") + a real tekken.json — the layout of an official mistralai Mistral-Small download, and the standalone-encoder layout these loaders explicitly accept — AutoTokenizer succeeds and returns MistralCommonBackend (mistral-common is a hard dependency, so this holds in every standard install).

That object BPE-encodes [SYSTEM_PROMPT]/[INST] as literal text instead of splicing them as single Tekken special ids. Verified against this module's own _TekkenRawTextAdapter on the same prompt: reference ends ..., 1, 17, 33153, 5117, 18, 3, 1097, 7990, 4 (markers = ids 17/18/3/4); MistralCommonBackend produces 1, 1091, 101289, 58343, ... (literal [, SY, STEM…). Every structural token is off-distribution, the encode "works", and generations silently degrade. Before this commit, that directory fell through to rung 4 and got the correct BFL processor whenever online/cached — so this regresses the online case too, and offline it trades an actionable RuntimeError for silent corruption.

Your ladder tests missed it because the fixture's tekken.json is a small fake that raises during parse; a real one parses fine and short-circuits.

Suggested fix: the ideal handling for a root tekken.json already exists in this file — wrap it with _TekkenRawTextAdapter (same as the embedded-tekken rung) before ever handing the directory to AutoTokenizer; and/or treat a mistral-common-backed result from AutoTokenizer as "not usable, keep falling". Either way, please add a fixture with a structurally valid tekken model file so the test exercises the success path, not just the raise path.

Non-blockers

1. The [dev]-side justification of the guard split overclaims, and the [dev] half is UI-dead. "buildFLUXGraph falls back to any FLUX.2 diffusers pipeline when only the VAE is needed" is true only of the Klein branch. The dev branch selects from dev-only pipelines (selectFlux2DevDiffusersModels) and readiness requires a dev diffusers source, so the config your test_dev_loader_extracts_vae_from_a_klein_pipeline proves the backend accepts (dev GGUF + standalone Mistral encoder + Klein diffusers as the only VAE source) still shows a disabled Invoke button. Safe direction, but either relax the dev builder/readiness VAE-sourcing to mirror Klein or trim the docstrings to the Klein case.
2. _validate_encoder_source in the Klein loader only rejects Dev, not Klein-family mismatches, while its comment claims the opaque-matmul class is closed. Klein 4B GGUF main + Klein 9B diffusers as qwen3_source_model (workflow editor) passes the guard and hits exactly that matmul error; the standalone-encoder path (_validate_qwen3_encoder_variant) and the frontend both enforce the family match — only the source path doesn't. Pre-existing, but worth closing while you're here: check the Klein→Qwen3 family map in the source path too. Related asymmetry: the dev guard rejects any non-dev variant, the Klein guard only dev — a future third FLUX.2 variant would be rejected on one side and silently accepted on the other.

Requesting changes for the two tokenizer-ladder blockers.